import os
import struct
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from    alive_progress import alive_bar

# 读取数据
def load_mnist(path, kind='train'):
    """Load MNIST data from `path`"""
    labels_path = path + '/' + kind + '-labels.idx1-ubyte'
    images_path = path + '/' + kind + '-images.idx3-ubyte'

    with open(labels_path, 'rb') as lbpath:
        magic, n = struct.unpack('>II',
                                 lbpath.read(8))
        labels = np.fromfile(lbpath,
                             dtype=np.uint8)

    with open(images_path, 'rb') as imgpath:
        magic, num, rows, cols = struct.unpack('>IIII',
                                               imgpath.read(16))
        images = np.fromfile(imgpath,
                             dtype=np.uint8).reshape(len(labels), 784)

    return images, labels

def visual_mnist(images, labels):
    # 查看mnist数据集的图片
    fig, ax = plt.subplots(
        nrows=2,
        ncols=5,
        sharex=True,
        sharey=True, )

    ax = ax.flatten()
    for i in range(10):
        img = images[labels == i][0].reshape(28, 28)
        ax[i].imshow(img, cmap='Greys', interpolation='nearest')

    ax[0].set_xticks([])
    ax[0].set_yticks([])
    plt.tight_layout()
    plt.show()


# 模型定义
class NaiveBayes():
    def __init__(self,x,y):
        _,n = np.shape(x)
        self.y_unique = np.unique(y)
        l   = len(np.unique(self.y_unique))
        self.xprob = np.zeros([l,n]) # 训练集各个特征在某一类下的方差
        self.yprob = np.zeros([l])   # 训练集y的分布
        # self.fit_ML(x,y)
        self.fit_bayes(x,y,1)

    def fit_ML(self,x,y):

        data_slice = [ np.squeeze(x[np.argwhere(y==i)],axis=1) for i in self.y_unique ]
        for i in range(len(self.y_unique)):
            self.xprob[i] = np.sum(data_slice[i],axis=0) / np.shape(data_slice[i])[0]
            self.yprob[i] = np.shape(data_slice[i])[0] / len(y)

        return 'Complete!'

    def fit_bayes(self,x,y,lamb):

        data_slice = [ np.squeeze(x[np.argwhere(y==i)],axis=1) for i in self.y_unique ]
        for i in range(len(self.y_unique)):
            self.xprob[i] = (np.sum(data_slice[i],axis=0) + lamb) / (np.shape(data_slice[i])[0] + 2*lamb)
            self.yprob[i] = np.shape(data_slice[i])[0] / len(y)

        return 'Complete!'

    def score(self,x,y):

        y_pred = []
        with alive_bar(len(y)) as bar: # declare your expected total

            for i in range(len(y)):
                y_pred.append(self.pred(x[i]))
                bar()

        total_num = len(y)
        corre_num = np.sum( (y_pred == y).astype(int) )
        acc = corre_num / total_num

        return acc

    def pred(self,x):

        prob = self.cal_prob(x)
        final_prob = np.ones([np.shape(prob)[0]])
        for i in range(np.shape(prob)[1]):
            final_prob = final_prob * prob[:,i]
        final_prob = final_prob * self.yprob
        max_index = np.argmax(final_prob)

        return self.y_unique[max_index]

    def cal_prob(self,x):
        
        prob = x*self.xprob + (1-x)*(1-self.xprob)

        return prob


# 主函数
def main():

    # 加载数据，打印基本信息
    images_train, labels_train = load_mnist('../Datasets/mnist')
    print('Images_train Shape:{:<25s} Labels_train Shape:{:<25s}'
            .format(str(images_train.shape),str(labels_train.shape)))
    print('Images_train type :{:<25s} Labels_train type :{:<25s}'
            .format(str(type(images_train)),str(type(labels_train))))
    images_test, labels_test = load_mnist('../Datasets/mnist','t10k')
    print('Images_test Shape :{:<25s}  Labels_test Shape:{:<25s}'
            .format(str(images_test.shape),str(labels_test.shape)))
    print('Images_test type  :{:<25s}  Labels_test type :{:<25s}'
            .format(str(type(images_test)),str(type(labels_test))))

    # 特征数据二值化
    Threshold = 127
    images_train = (images_train > Threshold).astype(int)
    images_test  = (images_test  > Threshold).astype(int)

    # 创建模型与训练
    bayes = NaiveBayes(images_train,labels_train)

    # 测试
    ACC = bayes.score(images_test,labels_test)
    print('ACC:',ACC)

if __name__ == '__main__':
    main()